import os
import sys

from IPython.display import Markdown
from IPython.display import display
from openai import OpenAI
sys.path.append('/JiuTian/swarm/swarm-main/swarm')
from swarm import Swarm,Agent

def run_multi_turn_conversation(
    openai_client,
    starting_agent,
    context_variables=None,
    debug=False
) -> None:
    # 创建 Swarm 客户端
    client = Swarm(openai_client)
    display(Markdown("## 开启Swarm对话 🐝").data)

    # 初始化消息列表
    messages = []
    agent = starting_agent  # 初始智能体

    while True:
        # 从用户获取输入
        user_input = input("User: ")
        if user_input.lower() in ["exit", "quit"]:
            display(Markdown("### Conversation Ended").data)
            break

        # 将用户输入添加到消息列表中
        messages.append({"role": "user", "content": user_input})

        # 运行 Swarm 客户端，智能体处理消息
        response = client.run(
            agent=agent,
            messages=messages,
            context_variables=context_variables or {},
            debug=debug,
        )

        # 使用 display(Markdown) 打印用户消息和智能体回复
        for message in response.messages:
            if message['role'] == 'user':
                display(Markdown(f"**User**: {message['content']}").data)
            elif message['role'] == 'assistant':
                display(Markdown(f"**{message['sender']}**: {message['content']}").data)

        # 更新消息和当前的智能体
        messages.extend(response.messages)
        agent = response.agent

#测试基本功能
client = OpenAI(api_key=os.environ["DEEPSEEK_API_KEY"], base_url=os.environ["DEEPSEEK_BASE_URL"])

def transfer_to_agent_b():
    return agent_b


model_name = "deepseek-chat"
#测试多智能体功能
swarm_client = Swarm(client)
agent_a = Agent(name = "Agent A",instructions= "你是一个乐于助人的助手",functions=[transfer_to_agent_b],model=model_name)
agent_b = Agent(name = "Agent B",instructions= "只使用俳句回答。",model=model_name)

run_multi_turn_conversation(openai_client = client,
                            starting_agent = agent_a)